import os
import random
import shutil
import time


# 划分样本函数
def Division(path, s_train_d, s_test_d, s_vali_d):
    image_or_list = os.listdir(path)  # 以列表形式获取原始样本
    save_list = []  # 空列表，用于选取样本的储存
    for i in range(sample_n):  # 把样品存储到列表save_list的过程
        save_list.append(image_or_list[i])
    tr_n = int(sample_n * tr_rate)
    tr_sam = random.sample(save_list, tr_n)  # 从save_list中随机获取3/4比例的图像，以列表的形式返回
    # 减去训练集样品剩下re集的样品
    re_sam = list(set(save_list) - set(tr_sam))
    #re集的训练集样品数目
    te_n = int(len(re_sam) * test_rate)
    te_sam = random.sample(re_sam, te_n)
    #re集的验证集样品数目
    vali_sam = list(set(re_sam) - set(te_sam))

    # 复制图像到目标文件夹
    for j in range(sample_n):
        name = image_or_list[j]
        if name in tr_sam:
            shutil.copy(os.path.join(path, name), os.path.join(s_train_d, name))
        elif name in te_sam:
            shutil.copy(os.path.join(path, name), os.path.join(s_test_d, name))
        else:
            shutil.copy(os.path.join(path, name), os.path.join(s_vali_d, name))


# 划分读入训练、测试、验证集数据
if __name__ == '__main__':
    time_start = time.time()

    # 各集比例
    tr_rate = 0.5  # 训练集比例1/2
    test_rate = 0.5  # 测试集占re集比例
    sample_n = 10000  # 戴口罩、不戴口罩样本选取数量

    # 原始数据集路径
    or_facemask_path = r"D:\s2096\CelebA人脸数据集\facemask_align_celeba\facemask_align_celeba"
    or_nomask_path = r"D:\s2096\CelebA人脸数据集\img_align_celeba"

    # 保存路径
    # 训练集路径
    s_tr_d1 = './data/train/mask/'
    s_tr_d2 = './data/train/no_mask/'
    s_te_d1 = './data/test/mask/'
    s_te_d2 = './data/test/no_mask/'
    s_vali_d1 = './data/validation/mask/'
    s_vali_d2 = './data/validation/no_mask/'

    # 数据集类别及数量
    Division(or_facemask_path, s_tr_d1, s_te_d1, s_vali_d1)
    Division(or_nomask_path, s_tr_d2, s_te_d2, s_vali_d2)

    print('划分完毕!')
    time_end = time.time()
    print('训练集、测试集、验证集划分共耗时%s!' % (time_end - time_start))



# 确实文件目录

b_dir = './dataset'
os.mkdir(b_dir)


# 训练图像的目录
tr_dir = os.path.join(b_dir, 'train')
os.mkdir(tr_dir)
# 有口罩的图片的训练资料的目录
tr_mask_dir = os.path.join(tr_dir, 'mask')
os.mkdir(tr_mask_dir)
# 没有口罩的图片的训练资料的目录
tr_no_mask_dir = os.path.join(tr_dir, 'no_mask')
os.mkdir(tr_no_mask_dir)


# 验证图像的目录
vali_dir = os.path.join(b_dir, 'validation')
os.mkdir(vali_dir)
# 有口罩的图片的验证集目录
vali_mask_dir = os.path.join(vali_dir, 'mask')
os.mkdir(vali_mask_dir)
# 没有口罩的图片的验证集目录
vali_no_mask_dir = os.path.join(vali_dir, 'no_mask')
os.mkdir(vali_no_mask_dir)


# 测试图像的目录
test_dir = os.path.join(b_dir, 'test')
os.mkdir(test_dir)
# 有口罩的图片的测试数据集目录
test_mask_dir = os.path.join(test_dir, 'mask')
os.mkdir(test_mask_dir)
# 没有口罩的图片的测试数据集目录
test_no_mask_dir = os.path.join(test_dir, 'no_mask')
os.mkdir(test_no_mask_dir)
